from bzrc import BZRC, Command, ShootCommand
from PriorityQueue import PriorityQueue
import sys, math, time, optparse, random, pprint

KP = 1
KD = 2
attractMaxDist = 100
repulseMaxDist = 80
tangentMaxDist = 30
attractWeight = .1
repulseWeight = .3
tangentWeight = 50
minAttract = 10
maxVel = 25
coneWidth = math.pi / 8

NORMAL_COST = 1
MAX_COST = 15
PEN_DIST = MAX_COST - 1
NODE_DIST = 25
NODE_SKIP = 3
SNIPE_DIST = 350

class Agent(object):

	def __init__(self, bzrc):
		self.bzrc = bzrc
		self.constants = self.bzrc.get_constants()
		self.bases = self.bzrc.get_bases()
		self.myBase = [base for base in self.bases if base.color == self.constants['team']][0]
		self.commands = []
		self.decoyTarget = [250,150]
		self.state = ["escapeMaze", "escapeMaze"]
		self.pfield = PField(self.bzrc)
		enemyBase = [base for base in self.bases if base.color == 'green'][0]
		self.enemyBase = self.pfield.find_center(enemyBase.coords)
		self.lastTarget = [0, 0]
		self.lastAng = [0, 0]
		self.targetNode = [0, 0]
		baseCenter = self.pfield.find_center(self.myBase.coords)
		search = MazeSearch(bzrc, int(baseCenter[0]), int(baseCenter[1]), int(self.constants['worldsize']))
		self.path = search.search()
		self.flippedPath = False

	def tick(self):
		'''Some time has passed; decide what to do next'''
		# Get information from the BZRC server
		global myTeamColor
		mytanks, othertanks, flags, shots = self.bzrc.get_lots_o_stuff()
		self.mytanks = mytanks
		self.flags = flags
		self.greenFlag = [flag for flag in flags if flag.color == 'green'][0]
		self.shots = shots
		self.enemies = [tank for tank in othertanks if tank.color ==
				'green']

		# Reset my set of commands (we don't want to run old commands)
		self.commands = []

		# Decide what to do with each of my tanks
		self.snipe(mytanks[0])
		self.decoy(mytanks[1])

		# Send the commands to the server
		results = self.bzrc.do_commands(self.commands)

	def snipe(self, bot):
		state = self.state[0]
		if state == "escapeMaze":
			self.escape_maze(bot)
		elif state == "primary":
			self.be_sniper(bot)
		elif state == "enemiesDead":
			self.get_flag(bot)
		elif state == "flagGet":
			self.return_home(bot)
			
	def decoy(self, bot):
		state = self.state[1]
		if state == "escapeMaze":
			self.escape_maze(bot)
		elif state == "primary":
			self.be_decoy(bot)
		elif state == "enemiesDead":
			self.get_flag(bot)
		elif state == "flagGet":
			self.return_home(bot)

	def be_decoy(self, bot):
		self.pfield.update()
		distance = math.hypot(self.decoyTarget[0] - bot.x, self.decoyTarget[1] - bot.y)
		if distance < 5:
			self.decoyTarget[1] *= -1
		(desiredVel, desiredAngle) = self.pfield.calc_field(bot.x, bot.y, self.decoyTarget[0], self.decoyTarget[1], True)
		delAngle = self.normalize_angle(desiredAngle - bot.angle)
		command = Command(bot.index, desiredVel, delAngle)
		self.commands.append(command)


	def be_sniper(self, bot):
		if self.mytanks[1].x < bot.x: #Don't get ahead of the decoy
			command = Command(bot.index, .1, 0)
			self.commands.append(command)
			return
		distance = math.hypot(self.enemyBase[0] - bot.x, self.enemyBase[1] - bot.y)
		if distance > SNIPE_DIST:
			vel, ang = self.pd(bot, (self.enemyBase[0], self.enemyBase[1]))
			command = Command(bot.index, vel, ang)
			self.commands.append(command)
		else:
			if self.enemies[1].status == "alive":
				tank = self.enemies[1]
			elif self.enemies[0].status == "alive":
				tank = self.enemies[0]
			else:
				print "Both enemies dead!"
				self.state = ["enemiesDead", "enemiesDead"]
				return
			desiredAngle = math.atan2(tank.y - bot.y, tank.x - bot.x)
			delAng = self.normalize_angle(desiredAngle - bot.angle)
			command = Command(bot.index, 0, delAng)
			self.commands.append(command)
			if delAng < 3:
				command = ShootCommand(bot.index)
				self.commands.append(command)

	def get_flag(self, bot):
		if self.greenFlag.poss_color == 'red':
			self.state[bot.index] = "flagGet"
			return
		vel, ang = self.pd(bot, (self.greenFlag.x, self.greenFlag.y))
		command = Command(bot.index, vel, ang)
		self.commands.append(command)

	def escape_maze(self, bot):
		nodeCoords = self.path[self.targetNode[bot.index]]
		vel, ang = self.pd(bot, nodeCoords)
		command = Command(bot.index, vel, ang)
		self.commands.append(command)
		dist = math.hypot(bot.x - nodeCoords[0], bot.y - nodeCoords[1])
		if dist < NODE_DIST:
			self.targetNode[bot.index] += NODE_SKIP
			if self.targetNode[bot.index] == len(self.path) - 1:
				self.state[bot.index] = "primary"
			elif self.targetNode[bot.index] > len(self.path):
				self.targetNode[bot.index] = len(self.path)-1
			#if bot.index == 0:
				#print "Bot ", bot.index, " now aims at ", self.path[self.targetNode[bot.index]]
		
	def return_home(self, bot):
		if not self.flippedPath:
			self.path.reverse()
			self.targetNode[0] = 0
			self.targetNode[1] = 0
			self.flippedPath = True
		self.escape_maze(bot)
		
		
	def pd(self, bot, targetPos):
		targetAng = math.atan2(targetPos[1] - bot.y, targetPos[0] - bot.x)
		botAng = self.normalize_angle(bot.angle)
		delA = (targetAng - botAng) - (self.lastTarget[bot.index] - self.lastAng[bot.index])
		a = KP * (targetAng - botAng) + KD * delA
		self.lastTarget[bot.index] = targetAng
		self.lastAng[bot.index] = botAng
		a = self.normalize_angle(a)
		return 1 - math.fabs(a/math.pi), a

	"""
	def pd(self, bot, targetPos):
		delY = (targetPos[1] - bot.y) - (self.lastTarget[bot.index][1] - self.lastPos[bot.index][1])
		delX = (targetPos[0] - bot.x) - (self.lastTarget[bot.index][0] - self.lastPos[bot.index][0])
		aY = KP * (targetPos[1] - bot.y) + KD * delY
		aX = KP * (targetPos[0] - bot.x) + KD * delX
		self.lastTarget[bot.index] = targetPos
		desiredAngle = math.atan2(aY, aX)
		angle = 2 * self.normalize_angle(desiredAngle - bot.angle)
		accel = math.hypot(aX, aY)
		print "PD Angle: ", angle/math.pi, " and accel: ", accel
		return accel, angle
	"""

	def behave(self, bot):
		if (bot.flag != '-'):
			ignoreFriends = True
			targetX, targetY = self.pfield.find_center(self.myBase.coords)
		else :
			targetX, targetY = self.closest_flag(bot)
			ignoreFriends = False
		(desiredVel, desiredAngle) = self.pfield.calc_field(bot.x, bot.y, targetX, targetY, ignoreFriends)
		delAngle = self.normalize_angle(desiredAngle - bot.angle)
		command = Command(bot.index, desiredVel, 2 * delAngle)
		self.check_los(bot)
		self.commands.append(command)

	def normalize_angle(self, angle):
		'''Make any angle be between +/- pi.'''
		angle -= 2 * math.pi * int (angle / (2 * math.pi))
		if angle <= -math.pi:
			angle += 2 * math.pi
		elif angle > math.pi:
			angle -= 2 * math.pi
		return angle
		
	def check_los(self, bot):
		fire = False
		minDist = () # Infinity
		global coneWidth
		for tank in self.othertanks:
			angle = math.atan2(tank.x - bot.x, tank.y - bot.y)
			minDist = min(math.hypot(tank.x - bot.x, tank.y - bot.y), minDist)
			if self.normalize_angle(angle - bot.angle) < coneWidth:
				fire = True
				break
		if fire:
			for tank in self.mytanks: # Check for friendly fire
				if tank.index == bot.index:
					continue
				angle = math.atan2(tank.x - bot.x, tank.y - bot.y)
				ffDist = math.hypot(tank.x - bot.x, tank.y - bot.y)
				if angle - bot.angle < coneWidth and ffDist < minDist:
					fire = False
					break
		if fire: # Clear to shoot at enemy
			command = ShootCommand(bot.index)
			self.commands.append(command)


class PField(object):
	
	def __init__(self, bzrc):
		self.bzrc = bzrc
		self.obstacles = bzrc.get_obstacles()
		
	# Helper function for adding vectors
	def add_vectors(self, vector1, vector2):
		# Get the x and y components
		x1 = vector1["magnitude"] * math.cos(vector1["angle"])
		y1 = vector1["magnitude"] * math.sin(vector1["angle"])
		x2 = vector2["magnitude"] * math.cos(vector2["angle"])
		y2 = vector2["magnitude"] * math.sin(vector2["angle"])

		xtot = x1 + x2
		ytot = y1 + y2
		newMag = math.sqrt(xtot * xtot + ytot * ytot)
		newAng = math.atan2(ytot, xtot)
		newVect = {"magnitude" : newMag, "angle" : newAng}
		return newVect
		
	# Calculates the total strength of the field at a given point.
	def calc_field(self, botX, botY, targetX, targetY, ignore_friends=False):
		otherTanks = self.otherTanks
		field = {"angle" : 0, "magnitude" : 0}

		
		field = self.add_vectors(field, self.attract_field(botX, botY, targetX, targetY))

		for tank in otherTanks:
			field = self.add_vectors(field, self.repulse_field(botX, botY, tank.x, tank.y))
		
		if not ignore_friends:
			for tank in self.myTanks:
				field = self.add_vectors(field, self.repulse_field(botX, botY, tank.x, tank.y))

		for obstacle in self.obstacles:
			field = self.add_vectors(field, self.tangential_field(botX, botY, obstacle))

		return self.normalize_vector(field)
				
	
	# Calculates a repulsive field given a point we want to avoid
	def repulse_field(self, botX, botY, startX, startY):
		distance = math.hypot(startX - botX, startY - botY)
		newMag = 0
		newAng = math.atan2((startY - botY), (startX - botX)) + math.pi
		
		if distance == 0:
			distance = 1

		if distance < repulseMaxDist:
			newMag = repulseWeight * repulseMaxDist / distance

		newVect = {"magnitude" : newMag, "angle" : newAng}

		return newVect
		
	
	# Calculates an attractive field given a point we want to reach
	def attract_field(self, botX, botY, startX, startY):
		distance = math.hypot(startX - botX, startY - botY)
		newMag = 0
		newAng = math.atan2((startY - botY), (startX - botX))
		
		if distance == 0:
			distance = 1

		if distance < attractMaxDist:
			newMag = attractWeight * distance
		else:
			newMag = attractWeight * attractMaxDist
			
		if newMag < minAttract:
			newMag = minAttract

		newVect = {"magnitude" : newMag, "angle" : newAng}

		return newVect
	
	# Calculates a tangential field given an obstacle we want to avoid
	def tangential_field(self, botX, botY, obstacle):
		(xcenter, ycenter) = self.find_center(obstacle)
		
		# So now we can calculate field strength
		distance = math.hypot(xcenter - botX, ycenter - botY)
		estRadius = math.hypot(obstacle[1][0] - obstacle[0][0], obstacle[1][1] - obstacle[0][1]) / 2
			
		if distance == 0:
			distance = 1
		
		if distance > estRadius:
			distance = distance - estRadius
			
		newMag = 0
		newAng = math.atan2((ycenter - botY), (xcenter - botX)) + math.pi

		if distance < tangentMaxDist:
			newMag = tangentWeight * tangentMaxDist / distance
			
		newAng = newAng - math.pi / 2 - random.uniform(0, 0.1) # I want the tangential field to be slightly repulsive, as well.

		newVect = {"magnitude" : newMag, "angle" : newAng}

		return newVect
		
	# Helper function for finding the center of an obstacle
	def find_center(self, obstacle):
		return (sum([pair[0] for pair in obstacle])/len(obstacle), sum([pair[1] for pair in obstacle])/len(obstacle))
				
	# Helper function for normalizing a vector by maxVel
	def normalize_vector(self, vect):
		return (vect["magnitude"] * 2, vect["angle"])
		
	def update(self):
		self.otherTanks = self.bzrc.get_othertanks()
		self.myTanks = self.bzrc.get_mytanks()


class MazeSearch:
	def __init__(self, bzrc, curX, curY, size):
		nodes, goal, penalties = self.parse_grid(size, bzrc)
		self.nodes = nodes
		self.startPos = (curX, curY)
		self.flagPos = (goal[0],goal[1])
		self.lastNode = (curX, curY)
		self.frontier = PriorityQueue()
		self.frontier.enqueue(0, (curX, curY), None, 0) 
		self.path = []
		self.searched = {}
		self.newNodes = []
		self.penalties = penalties
		
	def search(self):
		exists = self.iter_search()
		if not exists:
			print "Error! Could not find path from ", self.startPos, " to ", self.flagPos
				
		finalPath = [exists]
		curNode = exists
		nextNode = self.searched[exists]
		while nextNode != None:
			finalPath.insert(0, nextNode)
			curNode = nextNode
			nextNode = self.searched[curNode]
		return finalPath
		
	def step(self, steps):
		toReturn = []
		for i in range(steps):
			if not self.frontier:
				break
			done = self.iter_step()
			toReturn += self.newNodes
			if done:
				break
		return toReturn
		
	def iter_search(self):
		while self.frontier.has_item():
			answer = self.iter_step()
			if answer:
				return answer
		return False
		
	def iter_step(self):
		priority, cost, curNode, parent = self.frontier.dequeue()
		if not curNode in self.searched:
			self.searched[curNode] = parent
		if (curNode[0], curNode[1]) == self.flagPos:
			return curNode
		self.newNodes = []
		self.check_node((curNode[0]+1, curNode[1]-1), cost, curNode)
		self.check_node((curNode[0]-1, curNode[1]-1), cost, curNode)
		self.check_node((curNode[0]+1, curNode[1]+1), cost, curNode)
		self.check_node((curNode[0]-1, curNode[1]+1), cost, curNode)
		self.check_node((curNode[0], curNode[1]-1), cost, curNode)
		self.check_node((curNode[0], curNode[1]+1), cost, curNode)
		self.check_node((curNode[0]+1, curNode[1]), cost, curNode)
		self.check_node((curNode[0]-1, curNode[1]), cost, curNode)
		return False
		
	def check_node(self, node, cost, parent):
		if node in self.nodes and self.nodes[node] and not node in self.searched:
			newCost = cost + self.calc_cost(node, parent)
			if not self.frontier.get_data_priority(node):
				self.frontier.enqueue(newCost + self.heuristic_cost(node), node, parent, newCost)
				self.newNodes.append(node)
			elif self.frontier.get_data_priority(node) > newCost + self.heuristic_cost(node):
				self.frontier.reprioritize(newCost + self.heuristic_cost(node), node, parent, newCost)
				if node in self.searched:
					self.searched[curNode] = parent
				
	def heuristic_cost(self, node):
		#Our heuristic is just straight-line distance
		return math.hypot(node[0] - self.flagPos[0], node[1] - self.flagPos[1])
			
	def calc_cost(self, node, parent):
		return self.penalties[node[0],node[1]] * math.hypot(node[0] - parent[0], node[1] - parent[1])
						
	
	def is_adjacent(self, node):
		if (node[0],node[1]+1) in self.nodes and (self.nodes[(node[0],node[1]+1)] == 0):
			return True
		if (node[0]+1,node[1]+1) in self.nodes and (self.nodes[(node[0]+1,node[1]+1)] == 0):
			return True
		if (node[0]-1,node[1]+1) in self.nodes and (self.nodes[(node[0]-1,node[1]+1)] == 0):
			return True
		if (node[0]-1,node[1]) in self.nodes and (self.nodes[(node[0]-1,node[1])] == 0):
			return True
		if (node[0]+1,node[1]+1) in self.nodes and (self.nodes[(node[0]+1,node[1]+1)] == 0):
			return True
		if (node[0]+1,node[1]-1) in self.nodes and (self.nodes[(node[0]+1,node[1]-1)] == 0):
			return True
		if (node[0],node[1]-1) in self.nodes and (self.nodes[(node[0],node[1]-1)] == 0):
			return True
		if (node[0]-1,node[1]-1) in self.nodes and (self.nodes[(node[0]-1,node[1]-1)] == 0):
			return True
		return False
		
	def clear(self):
		self.path = []
		self.frontier = PriorityQueue()
		self.frontier.enqueue(0, (self.startPos[0], self.startPos[1]), None, 0)
		self.searched = {}
		
	def parse_grid(self, size, bzrc):
		toReturn = {}
		penaltyGrid = {}
		
		#Get the obstacles
		obstacles = bzrc.get_obstacles()
		otherTanks = bzrc.get_othertanks()
		
		furthestRight = -size/2
		
		for i in range(0, size):
			for j in range(0, size):
				toReturn[(int(i-size/2),int(j-size/2))] = 1
				penaltyGrid[(int(i-size/2),int(j-size/2))] = NORMAL_COST
		
		#Iterate through the obstacles, drawing line in between them
		for obstacle in obstacles:
			for i, point in enumerate(obstacle):
				#figure out which direction to go
				if i+1 < len(obstacle):
					point2 = obstacle[i+1]
				else:
					point2 = obstacle[0]
					
				diffx = point2[0] - point[0]
				diffy = point2[1] - point[1]
				
				rge = {}
				
				if diffx < 0:
					rge = range(0, int(diffx), -1)
				elif diffx > 0:
					rge = range(0, int(diffx), 1)
					
				if rge:
					for x in rge:
						toReturn[int(point[0])+x, int(point[1])] = 0
						for y in range(-PEN_DIST, PEN_DIST, 1):
							penaltyGrid[int(point[0])+x, int(point[1])+y] = MAX_COST - math.fabs(y)						
				else:
					if diffy < 0:
						rge = range(0, int(diffy), -1)
					else:
						rge = range(0, int(diffy), 1)
					for y in rge:
						toReturn[int(point[0]), int(point[1])+y] = 0
						for x in range(-PEN_DIST, PEN_DIST, 1):
							penaltyGrid[int(point[0])+x, int(point[1])+y] = MAX_COST - math.fabs(x)	
						
				if point[0] > furthestRight:
					furthestRight = point[0]
					
		for i in range(int(furthestRight+1), int(size/2)):
			for j in range(0, size):
				toReturn[(int(i),int(j-size/2))] = 0
				
		goal = False
		yloc = -size/2+1
		
		while not goal:
			if toReturn[int(furthestRight), int(yloc)]:
				goal = (furthestRight, yloc + 25)
			yloc = yloc + 1
					
		return toReturn, goal, penaltyGrid


def main():
	# Process CLI arguments.
	parser = optparse.OptionParser(usage = "%prog [options]",
								   version = "%prog 0.1")
	parser.add_option("-o", "--host", type="string",dest="host",
			metavar="HOST",default='localhost',
			help="hostname for the server")
	parser.add_option("-p", "--port", type="int",dest="port",
			metavar="PORT",default=3907,
			help="port number for the server")

	(options, args) = parser.parse_args()

	# Connect.
	#bzrc = BZRC(options.host, int(options.port), debug=True)
	bzrc = BZRC(options.host, int(options.port))

	agent = Agent(bzrc)

	# Run the agent
	try:
		while True:
			agent.tick()
	except KeyboardInterrupt:
		print "Exiting due to keyboard interrupt."
		bzrc.close()



if __name__ == '__main__':
	main()
